# encoding:utf-8
import rados, rbd
import functools
import time
from tornado.concurrent import Future
from tornado.options import options


def unicode_to_string(f):
    '''对返回结果进行统一编码'''

    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        if args:
            arg_list = []
            for arg in args:
                if isinstance(arg, unicode): arg = arg.encode('utf-8')
                arg_list.append(arg)
            args = arg_list
        if kwargs:
            prams = {}
            for key in kwargs.keys():
                pram = kwargs.get(key)
                if isinstance(pram, unicode): pram = pram.encode('utf-8')
                prams[key] = pram
            kwargs = prams

        return f(*args, **kwargs)

    return wrapper


def wrapper_response(f):
    '''对接口点用后的返回值统一组装返回格式'''

    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        try:
            response = f(*args, **kwargs)
            if isinstance(response, (rbd.SnapIterator, rados.ObjectIterator, rados.XattrIterator), ):
                response_list = []
                for res in response:
                    if isinstance(res, rados.Object):
                        res = res.key
                        response_list.append(res)
                    if isinstance(res, tuple):
                        response_list.append({res[0]: res[1]})
                    if isinstance(res, dict):
                        response_list.append(res)
                response = {'status': 'OK', 'output': response_list}

            elif isinstance(response, bool):
                if response:
                    response = {'status': 'OK', 'output': ''}
                else:
                    response = {'status': 'ERROR', 'output': ''}
            elif isinstance(response, tuple):
                strftime = time.strftime('%Y-%m-%d %H:%M:%S', response[1])
                response = {'status': "OK", 'output': (int(response[0]), strftime,)}
            else:
                response = {'status': 'OK', 'output': response if response else ''}

        except Exception as e:
            response = {'status': 'ERROR', 'output': '{0}:{1}'.format(type(e), e.message)}
        return response

    return wrapper


class CephConnection(object):
    futures = {}
    ioctxs = {}
    rbds = {}

    def future(self, object_id):
        future = Future()
        self.futures[object_id] = future
        return future

    def __init__(self, config_file=None, key_file=None):
        config_file = config_file if config_file else options.ceph
        key_file = key_file if key_file else options.ceph_key
        self.cluster = rados.Rados(conffile=config_file, conf=dict(keyring=key_file))
        self.cluster.connect()

    @classmethod
    def instance(cls, config_file=None, key_file=None):
        if not hasattr(cls, '_instance'):
            cls._instance = cls(config_file, key_file)
        return cls._instance

    def __del__(self):
        for name, rbd in self.rbds:
            rbd.break_lock()
        for name, ioctx in self.ioctxs:
            ioctx.close()
        if hasattr(self, 'cluster'):
            self.cluster.shutdown()

    def open_rbd(self, name):
        if name not in self.rbds.keys():
            self.rbds[name] = rbd.Image(self.ioctx, name)
        return self.rbds[name]

    @unicode_to_string
    @wrapper_response
    def list_pool(self):
        """返回所有集群"""
        return self.cluster.list_pools()

    @unicode_to_string
    @wrapper_response
    def pool_list_objects(self, poolname):
        """获取一个池所有对象"""
        self.open_ioctx(poolname)
        return self.ioctx.list_objects()

    @unicode_to_string
    @wrapper_response
    def delete_pool_object(self, poolname, file_name):
        """删除一个对象"""
        self.open_ioctx(poolname)
        return self.ioctx.remove_object(file_name)

    @unicode_to_string
    @wrapper_response
    def pools(self):
        """获取池列表"""
        return self.cluster.list_pools()

    @unicode_to_string
    def async_write_full(self, pool, name, data, obj_id, offset=0, func=None):
        """异步写对象"""
        self.open_ioctx(pool)
        future = self.future(obj_id)
        self.ioctx.aio_write(name, data, offset=offset, oncomplete=func, onsafe=self.write_callback(obj_id))
        return future

    @unicode_to_string
    def async_append_full(self, pool, name, data, obj_id):
        """异步增加对象"""
        self.open_ioctx(pool)
        future = self.future(obj_id)
        self.ioctx.aio_append(name, data, onsafe=self.write_callback(obj_id))
        return future

    @unicode_to_string
    def write_callback(self, obj_id):
        me = self

        def _lambda(completion):
            if completion:
                me.futures.pop(obj_id).set_result(completion)
            else:
                me.futures.pop(obj_id).set_exception(ValueError)

        return _lambda

    @unicode_to_string
    def open_ioctx(self, pool):
        if pool not in self.ioctxs.keys():
            self.ioctxs[pool] = self.cluster.open_ioctx(pool)
        self.ioctx = self.ioctxs.get(pool)

    @unicode_to_string
    @wrapper_response
    def object_info(self, pool, name):
        """获取对象属性"""
        self.open_ioctx(pool)
        return self.ioctx.stat(name)

    @unicode_to_string
    @wrapper_response
    def pool_info(self, pool):
        """获取对象属性"""
        self.open_ioctx(pool)
        return self.ioctx.get_stats()

    @unicode_to_string
    def async_read_full(self, pool, name, obj_id, length=0, offset=0):
        """异步读对象"""
        self.open_ioctx(pool)
        future = self.future(obj_id)
        self.ioctx.aio_read(name, length, offset, self.read_callback(obj_id))
        return future

    @unicode_to_string
    @wrapper_response
    def obj_set_xattr(self, pool, name, key, val):
        """设置对象一个meta"""
        self.open_ioctx(pool)
        return self.ioctx.set_xattr(name, key, val)

    @unicode_to_string
    @wrapper_response
    def obj_get_xattrs(self, pool, name):
        """获取对象meta"""
        self.open_ioctx(pool)
        return self.ioctx.get_xattrs(name)

    @unicode_to_string
    @wrapper_response
    def obj_del_xattr(self, pool, name, key):
        """删除对象meta"""
        self.open_ioctx(pool)
        return self.ioctx.rm_xattr(name, key)

    @unicode_to_string
    @wrapper_response
    def obj_list(self, pool):
        """获取对象列表"""
        self.open_ioctx(pool)
        return self.ioctx.list_objects()

    @unicode_to_string
    def read_callback(self, obj_id):
        me = self

        def _lambda(completion, data_read):
            if data_read and completion:
                me.futures.pop(obj_id).set_result(data_read)
            else:
                me.futures.pop(obj_id).set_exception(ValueError)

        return _lambda

    @unicode_to_string
    def async_append(self, poolname, name, data, obj_id):
        self.open_ioctx(poolname)
        future = self.future(obj_id)
        self.ioctx.aio_append(name, data, self.write_callback(obj_id))
        return future

    @unicode_to_string
    def read(self, pool, name):
        """读取一个块"""
        self.open_ioctx(pool)
        return self.ioctx.aio_read(name)

    @unicode_to_string
    @wrapper_response
    def create_rbd_snap(self, pool, block_story_name, name):
        """创建块的快照"""
        self.open_ioctx(pool)
        img = rbd.Image(self.ioctx, block_story_name)
        return img.create_snap(name)

    @unicode_to_string
    @wrapper_response
    def list_rbd_snap(self, pool, name):
        """获取块的快照列表"""
        self.open_ioctx(pool)
        img = rbd.Image(self.ioctx, name)
        return img.list_snaps()

    @unicode_to_string
    @wrapper_response
    def delete_rbd_snap(self, pool, block_story_name, name):
        """删除一个块的快照"""
        self.open_ioctx(pool)
        img = rbd.Image(self.ioctx, block_story_name, )
        return img.remove_snap(name)

    @unicode_to_string
    @wrapper_response
    def rollback_rbd_to_snap(self, pool, block_story_name, name):
        """回滚一个块到一个快照"""
        self.open_ioctx(pool)
        img = rbd.Image(self.ioctx, block_story_name, )
        return img.rollback_to_snap(name)

    @unicode_to_string
    @wrapper_response
    def protect_rbd_snap(self, pool, block_story_name, snap_name):
        """设置快照保护"""
        self.open_ioctx(pool)
        img = rbd.Image(self.ioctx, block_story_name)
        return img.protect_snap(snap_name)

    @unicode_to_string
    @wrapper_response
    def unprotect_rbd_snap(self, pool, block_story_name, snap_name):
        """取消快照保护"""
        self.open_ioctx(pool)
        img = rbd.Image(self.ioctx, block_story_name)
        return img.unprotect_snap(snap_name)

    @unicode_to_string
    @wrapper_response
    def rbd_is_protect_snap(self, pool, block_story_name, snap_name):
        """判断一个块是否保护"""
        self.open_ioctx(pool)
        img = rbd.Image(self.ioctx, block_story_name)
        return img.is_protected_snap(snap_name)

    @unicode_to_string
    @wrapper_response
    def list_rbd(self, pool):
        """获取块列表"""
        self.open_ioctx(pool)
        img = rbd.RBD()
        return img.list(self.ioctx)

    @unicode_to_string
    @wrapper_response
    def create_rbd(self, pool, name, size, features=1, stripe_unit=0,
                   stripe_count=0):
        """创建一个块"""
        self.open_ioctx(pool)
        rbd_info = rbd.RBD()
        response = rbd_info.create(self.ioctx, name, size, old_format=False, features=features, stripe_unit=stripe_unit,
                                   stripe_count=stripe_count)
        return response

    @unicode_to_string
    def write_rbd(self, pool, name, data, offset=0):
        """写块文件"""
        self.open_ioctx(pool)
        self.open_rbd(name)
        return self.rbd.write(data, offset=offset)

    @unicode_to_string
    @wrapper_response
    def delete_rbd(self, pool, name):
        """删除一个块"""
        self.open_ioctx(pool)
        return rbd.RBD().remove(self.ioctx, name)

    @unicode_to_string
    @wrapper_response
    def clone_rbd(self, pool, blockname, snapname, clone_pool, clone_blockname, ):
        """克隆一个块到新块"""
        self.open_ioctx(pool)
        block = rbd.RBD()
        c_ioctx = self.cluster.open_ioctx(clone_pool)
        clone = block.clone(self.ioctx, blockname, snapname, c_ioctx, clone_blockname, features=1, order=None)
        c_ioctx.close()
        return clone

    @unicode_to_string
    @wrapper_response
    def flatten_rbd(self, pool, name):
        """完全拷贝块文件"""
        self.open_ioctx(pool)
        block = rbd.Image(self.ioctx, name)
        return block.flatten()

    @unicode_to_string
    @wrapper_response
    def rbd_info(self, pool, name):
        """获取块文件信息"""
        self.open_ioctx(pool)
        block = rbd.Image(self.ioctx, name)
        return block.stat()

    @unicode_to_string
    @wrapper_response
    def resize_rbd(self, pool, name, size):
        """重置rbd大小"""
        self.open_ioctx(pool)
        block = rbd.Image(self.ioctx, name)
        return block.resize(size)


if __name__ == "__main__":
    import os

    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
    ceph = CephConnection('%s%s' % (BASE_DIR, '/../ceph.conf',), key_file='%s%s' % (BASE_DIR, '/../ceph.keyring',))
    # size = 4 * 1024 ** 3
    # ceph.rbd_create()
    # conf_ = '%s%s' % (BASE_DIR, '/../ceph.conf',)
    # print conf_
    # cluster = rados.Rados(conffile=conf_)
    # cluster.connect()
    # ioctx = cluster.open_ioctx('zh')
    # rbd_inst = rbd.RBD()
    size = 4 * 1024 ** 3  # 4 GiB
    ceph.create_rbd('zh', 'myimage', size)
